import os
import cv2
import numpy as np
import tifffile as tiff

# 设置文件夹路径
gt_folder = "C:/Users/SONGSHITAO/Desktop/gt"
pred_folder = "C:/Users/SONGSHITAO/Desktop/unet-73"
output_folder = "C:/Users/SONGSHITAO/Desktop/unet"

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)

# 获取 GT 和预测结果文件名（取交集，保证只处理匹配的文件）
gt_files = set(os.listdir(gt_folder))
pred_files = set(os.listdir(pred_folder))
common_files = sorted(gt_files & pred_files)  # 取交集，保证文件名相同

print(f"找到 {len(common_files)} 个匹配文件，开始处理...")

# 遍历匹配的文件
for filename in common_files:
    gt_path = os.path.join(gt_folder, filename)
    pred_path = os.path.join(pred_folder, filename)
    output_path = os.path.join(output_folder, filename)

    # **读取 GT 并归一化**
    gt = tiff.imread(gt_path).astype(np.uint8)
    gt = (gt // 255) * 255  # 确保 GT 只有 0 和 255

    # **读取推理结果**
    pred = tiff.imread(pred_path)

    # **转换推理结果为 0 和 255**
    if pred.ndim == 3:  # 如果是 RGB 图，先转换为灰度
        pred_gray = cv2.cvtColor(pred, cv2.COLOR_RGB2GRAY)
    else:
        pred_gray = pred

    # **二值化，确保前景=255，背景=0**
    _, pred_bin = cv2.threshold(pred_gray, 127, 255, cv2.THRESH_BINARY)

    # **确保形状匹配**
    if gt.shape != pred_bin.shape:
        print(f"警告: {filename} 形状不匹配，跳过处理 ({gt.shape} vs {pred_bin.shape})")
        continue

    # **创建彩色标注结果图像**
    h, w = gt.shape
    result = np.zeros((h, w, 3), dtype=np.uint8)  # 初始化为黑色

    # **标注错误区域**
    mask_building = (gt == 255) & (pred_bin == 0)  # GT=前景，预测=背景（错误）→ 红色
    mask_background = (gt == 0) & (pred_bin == 255)  # GT=背景，预测=前景（错误）→ 黄色

    result[mask_building] = [255, 0, 0]  # 红色 (前景误分类)
    result[mask_background] = [255, 255, 0]  # 黄色 (背景误分类)

    # **正确部分保持 GT 原样**
    mask_correct = (gt == pred_bin)
    result[mask_correct] = np.stack([gt[mask_correct]] * 3, axis=-1)  # 直接使用 GT 原值

    # **保存结果**
    tiff.imwrite(output_path, result)

print(f"✅ 处理完成！结果已保存到：{output_folder}")